import os
from starter_code.utils import load_case
import numpy as np

import cv2
from pre_process import data_augment
from config import cfg


# 自定义数据集
class DatasetGenerator:
    def __init__(self, cases_path):
        # 把所有图编号，放入一个列表
        self.cases_path = os.listdir(cases_path)
        self.imgs_dir = []
        for case_path in self.cases_path:
            volumes, segmentations = load_case(case_path)
            batch_size = volumes.shape[0]
            # num 是图片和标签在volume中的编号
            num = 0
            for i in range(batch_size):
                self.imgs_dir.append((num, case_path))
                num += 1

    def __getitem__(self, index):
        # case 是 1, num 是 0
        case_path = self.imgs_dir[index][1]
        num = self.imgs_dir[index][0]

        volumes, segmentations = load_case(case_path)
        img = volumes.get_fdata()[num]
        # img = np.transpose()
        label = segmentations.get_fdata()[num]

        label = np.reshape(label, (512, 512))
        img = np.reshape(img, (1, 512, 512))
        # label.shape[0], label.shape[1])
        label_zero = np.zeros((3, 512, 512))
        for cn in range(3):
            label_zero[cn, :, :] = np.array((label[:, :] == cn))
        # img = img /
        # label = label / 255.
        return img, label_zero

        # img_name = self.img_path[index]
        # img_item_path = os.path.join(self.root_dir, self.img_dir, img_name)
        # img = cv2.imread(img_item_path, cv2.IMREAD_COLOR)
        # label_name = self.label_path[index]
        # label_item_path = os.path.join(self.root_dir, self.label_dir,
        #                                label_name)
        # label = cv2.imread(label_item_path, cv2.IMREAD_GRAYSCALE)
        # data_augment(img, label)
        # img = np.transpose(img, (2, 0, 1))
        # label = np.reshape(label, (1, cfg.IMAGE_HEIGHT, cfg.IMAGE_WIDTH))
        # img = img / 255.
        # label = label / 255.
        # return img, label

    def __len__(self):
        return len(self.cases_path)
